#!/usr/bin/env python
import os
import subprocess
import numpy as np
import matplotlib.pyplot as plt
#subprocess.run(['conda', 'init'])
#subprocess.run(['conda', 'activate', 'deepmd'])
for i in range(0,1):
    for j in range(0,4):
        
        dft_dp_path = '/home/zxg/be-cu_dp_course/dpgen/run/deepmd_file/'
        init_fp_list = os.listdir(dft_dp_path)

        for file in init_fp_list:
            folder_list = os.path.join(dft_dp_path, file)
            
            energy_raw_file = os.path.join(folder_list, 'deepmd/energy.raw')
            print(f"now test for {file}")
            with open(energy_raw_file, "r") as energy_raw:
                lines = energy_raw.readlines()
            num_frame = str(len(lines))
            print(f"there are {num_frame} frame for dft")

            dp_file_path = f"/home/zxg/be-cu_dp_course/dpgen/run/iter.00000{i}/00.train/00{j}/frozen_model.pb"
            results = f"iter_{i}_{j}"
            print(f"now do the dp test for {results}")

            subprocess.run(['dp', 'test', '-m', dp_file_path, '-s', folder_list, '-n', num_frame, '-d', results], check=True)

            # 定义绘制散点图和对角线的函数
            def plot(ax, data, key, xlabel, ylabel, min_val, max_val):
                data_key = f'data_{key}'
                pred_key = f'pred_{key}'
                ax.scatter(data[data_key], data[pred_key], label=key, s=6)
                ax.legend()
                ax.set_xlabel(xlabel)
                ax.set_ylabel(ylabel)
                ax.set_xlim(min_val, max_val)
                ax.set_ylim(min_val, max_val)
                ax.plot([min_val, max_val], [min_val, max_val], 'r', lw=1)
            
            # 读取数据，并对e数据进行原子化处理
            type_raw_file = os.path.join(folder_list, 'deepmd/type.raw')
            with open(type_raw_file, "r") as type_raw:
                lines = type_raw.readlines()
            natom = len(lines)
            print(f"there are {natom} for dft")
            results_e = f"{results}.e_peratom.out"
            results_f = f"{results}.f.out"
            data_e = np.genfromtxt(results_e, names=["data_e", "pred_e"])
            data_f = np.genfromtxt(results_f, names=["data_fx", "data_fy", "data_fz", "pred_fx", "pred_fy", "pred_fz"])
            
            
            # 计算e和f的最小值和最大值
            data_e_stacked = np.column_stack((data_e['data_e'], data_e['pred_e']))
            data_f_stacked = np.column_stack((data_f['data_fx'], data_f['data_fy'], data_f['data_fz'], data_f['pred_fx'], data_f['pred_fy'], data_f['pred_fz']))
            
            min_val_e, max_val_e = np.min(data_e_stacked), np.max(data_e_stacked)
            min_val_f, max_val_f = np.min(data_f_stacked), np.max(data_f_stacked)
            
            # 绘制散点图并保存结果
            fig, axs = plt.subplots(1, 2, figsize=(12, 5))
            plot(axs[0], data_e, 'e', 'DFT energy (eV/atom)', 'DP energy (eV/atom)', min_val_e, max_val_e)
            for force_direction in ['fx', 'fy', 'fz']:
                plot(axs[1], data_f, force_direction, 'DFT force (eV/Å)', 'DP force (eV/Å)', min_val_f, max_val_f)
            plt.savefig(f'{results}.png')
            plt.clf()
#subprocess.run(['conda', 'deactivate'])